package com.springboot.distributed.transaction;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;

import javax.sql.DataSource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.jdbc.datasource.DistributedDataSourceTransactionObject;
import org.springframework.transaction.TransactionSystemException;
import org.springframework.transaction.support.DefaultTransactionStatus;

import com.springboot.common.util.SafeNullUtils;
import com.springboot.distributed.transaction.zookeeper.TransactionMessage;
import com.springboot.distributed.transaction.zookeeper.TransactionStatusEnum;

/**
 * 分布式事务数据源适配器
 *
 * @author supu
 * @date 2019-09-10 17:21
 **/
public class DistributedTransactionDataSourceAdapter implements TransactionStatusHandler, InitializingBean {
    private static final Logger LOGGER = LoggerFactory.getLogger(DistributedTransactionDataSourceAdapter.class);

    private DataSource dataSource;
    private DistributedTransactionHandler distributedTransactionHandler;
    private TransactionMessageHandler transactionMessageHandler;

    @Override
    public void afterPropertiesSet() throws Exception {
        transactionMessageHandler.init(this);
    }

    public void setDataSource(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    public void setDistributedTransactionHandler(DistributedTransactionHandler distributedTransactionHandler) {
        this.distributedTransactionHandler = distributedTransactionHandler;
    }

    public void setTransactionMessageHandler(TransactionMessageHandler transactionMessageHandler) {
        this.transactionMessageHandler = transactionMessageHandler;
    }

    /**
     * 完成之后清除操作
     *
     * @param txObject
     */
    public void doCleanupAfterCompletion(DistributedDataSourceTransactionObject txObject) {
        if (distributedTransactionHandler.isDistributedTransaction()) {
            distributedTransactionHandler.doCleanupAfterCompletion();
            return;
        }

        doRealCleanupAfterCompletion(txObject);
    }

    public void doCommit(DefaultTransactionStatus status) {
        if (distributedTransactionHandler.isDistributedTransaction()) {
            dealDefaultTransactionStatus(status, TransactionStatusEnum.COMMIT);
            return;
        }

        doRealCommit(status);
    }

    private void doRealCommit(DefaultTransactionStatus status) {
        DistributedDataSourceTransactionObject txObject = (DistributedDataSourceTransactionObject) status.getTransaction();
        Connection con = txObject.getConnectionHolder().getConnection();
        if (status.isDebug()) {
            LOGGER.debug("Committing JDBC transaction on Connection [" + con + "]");
        }
        try {
            con.commit();
        } catch (SQLException ex) {
            throw new TransactionSystemException("Could not commit JDBC transaction", ex);
        } finally {
            doRealCleanupAfterCompletion(txObject);
        }
    }

    public void doRollBack(DefaultTransactionStatus status) {
        if (distributedTransactionHandler.isDistributedTransaction()) {
            dealDefaultTransactionStatus(status, TransactionStatusEnum.ROLLBACK);
            return;
        }

        doRealRollBack(status);
    }

    private void dealDefaultTransactionStatus(DefaultTransactionStatus status, TransactionStatusEnum transactionStatusEnum) {
        boolean isHeader = distributedTransactionHandler.saveTransactionStatus(status);
        if (isHeader) {
            // 如果是事务入口，创建事务回滚状态消息保存到zk中
            transactionMessageHandler.sendTransactionMessage(
                    new TransactionMessage(distributedTransactionHandler.getUniqueKey(), transactionStatusEnum));
        }
    }

    private void doRealRollBack(DefaultTransactionStatus status) {
        DistributedDataSourceTransactionObject txObject = (DistributedDataSourceTransactionObject) status.getTransaction();
        Connection con = txObject.getConnectionHolder().getConnection();
        if (status.isDebug()) {
            LOGGER.debug("Rolling back JDBC transaction on Connection [" + con + "]");
        }
        try {
            con.rollback();
        } catch (SQLException ex) {
            throw new TransactionSystemException("Could not roll back JDBC transaction", ex);
        } finally {
            doRealCleanupAfterCompletion(txObject);
        }
    }

    /**
     * 完成之后真正的清除操作
     *
     * @param txObject
     */
    private void doRealCleanupAfterCompletion(DistributedDataSourceTransactionObject txObject) {
        // Reset connection.
        Connection con = txObject.getConnectionHolder().getConnection();
        try {
            if (txObject.isMustRestoreAutoCommit()) {
                con.setAutoCommit(true);
            }
            DataSourceUtils.resetConnectionAfterTransaction(con, txObject.getPreviousIsolationLevel());
        } catch (Throwable ex) {
            LOGGER.debug("Could not reset JDBC Connection after transaction", ex);
        }

        if (txObject.isNewConnectionHolder()) {
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("Releasing JDBC Connection [" + con + "] after transaction");
            }
            DataSourceUtils.releaseConnection(con, this.dataSource);
        }

        txObject.getConnectionHolder().clear();
    }

    @Override
    public boolean handlerTransaction(TransactionMessage transactionMessage) {
        List<DefaultTransactionStatus> transactionStatusList =
                distributedTransactionHandler.getDefaultTransactionStatusList(transactionMessage.getUniqueKey());
        LOGGER.info("transactionMessage:{} 对应的事务列表: {}", transactionMessage, transactionStatusList);
        SafeNullUtils.safeGetStream(transactionStatusList).forEach(transactionStatus -> {
            if (transactionMessage.isCommit()) {
                doRealCommit(transactionStatus);
            } else {
                doRollBack(transactionStatus);
            }
        });

        distributedTransactionHandler.clearDistributedTransaction(transactionMessage.getUniqueKey());

        return true;
    }


}
